import pickle as pkl
import pandas as pd
import random
import numpy as np
import time
from bidict import bidict

'''
Preprocess Articles
'''
df = pd.read_csv(".\\input\\articles.csv", sep=',', dtype=str)
artid_bidict, pcode_bidict = bidict(), bidict()
# start counting from 1, because 0 is reserved for mask value
artid_cnt, pcode_cnt = 0, 0
artid, pcode = [], []
for id, code in zip(df['article_id'], df['product_code']):
    if id not in artid_bidict:
        artid_cnt += 1
        artid_bidict[id] = artid_cnt
    if code not in pcode_bidict:
        pcode_cnt += 1
        pcode_bidict[code] = pcode_cnt
    artid.append(artid_bidict[id])
    pcode.append(pcode_bidict[code])
artid = np.array(artid, dtype='int')
pcode = np.array(pcode, dtype='int')
print("artid_cnt, pcode_cnt, artid.shape, pcode.shape", artid_cnt, pcode_cnt, artid.shape, pcode.shape)

with open('.\\input\\processed\\article_sparse.pkl', 'wb') as f:
    pkl.dump(artid, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pcode, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(artid_bidict, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pcode_bidict, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump((artid_cnt, pcode_cnt), f, pkl.HIGHEST_PROTOCOL)

print("Articles Preprocessed")

'''
Preprocess Customers
'''
df = pd.read_csv(".\\input\\customers.csv", sep=',')
df['FN'].fillna(0.0, inplace=True)
df['Active'].fillna(0.0, inplace=True)
df['club_member_status'].fillna('PRE-CREATE', inplace=True)
df['fashion_news_frequency'].fillna('NONE', inplace=True)
df['age'].fillna(df['age'].mean(), inplace=True)

for col in df.columns:
    detect_nan = df[col].isnull().values.any()
    print(col, detect_nan)

customer_bidict = bidict()
# start counting from 0
customer_cnt = 0
customer_id = []
for id in df['customer_id']:
    if id not in customer_bidict:
        customer_bidict[id] = customer_cnt
        customer_cnt += 1
    customer_id.append(customer_bidict[id])
customer_id = np.array(customer_id, dtype='int')
print("customer_cnt, customer_id.shape", customer_cnt, customer_id.shape)

FN = [0 if x < 0.5 else 1 for x in df['FN']]
Active = [0 if x < 0.5 else 1 for x in df['Active']]
club_member_status = [1 if x == 'ACTIVE' else 0 for x in df['club_member_status']]
fashion_news_frequency = [1 if x == 'Regularly' else 0 for x in df['fashion_news_frequency']]
age = [x for x in df['age']]

FN = np.array(FN, dtype='int')
Active = np.array(Active, dtype='int')
club_member_status = np.array(club_member_status, dtype='int')
fashion_news_frequency = np.array(fashion_news_frequency, dtype='int')
age = np.array(age, dtype='float')

print(FN.shape, Active.shape, club_member_status.shape, fashion_news_frequency.shape, age.shape)
FN, Active, club_member_status, fashion_news_frequency, age

with open('.\\input\\processed\\customer_feat.pkl', 'wb') as f:
    pkl.dump(customer_id, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(FN, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(Active, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(club_member_status, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(fashion_news_frequency, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(age, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(customer_cnt, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(customer_bidict, f, pkl.HIGHEST_PROTOCOL)

print("Customers preprocessed")


'''
Preprocess transactions
'''
df = pd.read_csv(".\\input\\transactions_train.csv", sep=',', dtype=str)
with open('.\\input\\processed\\article_sparse.pkl', 'rb') as f:
    artid = pkl.load(f)
    pcode = pkl.load(f)
    artid_bidict = pkl.load(f)
    pcode_bidict = pkl.load(f)
    artid_cnt, pcode_cnt = pkl.load(f)

with open('.\\input\\processed\\customer_feat.pkl', 'rb') as f:
    customer_id = pkl.load(f)
    FN = pkl.load(f)
    Active = pkl.load(f)
    club_member_status = pkl.load(f)
    fashion_news_frequency = pkl.load(f)
    age = pkl.load(f)
    customer_cnt = pkl.load(f)
    customer_bidict = pkl.load(f)

assert artid_cnt == artid.shape[0] == pcode.shape[0]
assert customer_cnt == customer_id.shape[0] == FN.shape[0] == Active.shape[0] == club_member_status.shape[0] \
    == fashion_news_frequency.shape[0] == age.shape[0]

df['customer_id'] = df['customer_id'].apply(lambda x: customer_bidict[x])
df['article_id'] = df['article_id'].apply(lambda x: artid_bidict[x])
df['price'] = df['price'].apply(lambda x: float(x))
df['sales_channel_id'] = df['sales_channel_id'].apply(lambda x: 0 if x == '1' else 1)

customer_idx = df['customer_id'].tolist()
article_idx = (df['article_id'] - 1).tolist()
pos_FN = FN[customer_idx]
pos_Active = Active[customer_idx]
pos_club_member_status = club_member_status[customer_idx]
pos_fashion_news_frequency = fashion_news_frequency[customer_idx]
pos_age = age[customer_idx]

pos_artid = artid[article_idx]
pos_pcode = pcode[article_idx]

print(pos_FN.shape, pos_Active.shape, pos_club_member_status.shape, pos_fashion_news_frequency.shape, pos_age.shape)
print(pos_artid.shape, pos_pcode.shape)

pos_customer_id = df['customer_id'].to_numpy()
print(pos_customer_id, pos_customer_id.shape)

df.to_csv('.\\input\\processed\\transactions_translated.csv')
with open('.\\input\\processed\\pos_feat.pkl', 'wb') as f:
    pkl.dump(pos_customer_id, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pos_FN, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pos_Active, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pos_club_member_status, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pos_fashion_news_frequency, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pos_age, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pos_artid, f, pkl.HIGHEST_PROTOCOL)
    pkl.dump(pos_pcode, f, pkl.HIGHEST_PROTOCOL)

print("Transactions processed")
